library(dplyr)
library(stringr)
library(ggplot2)
library(CNAqc)

set.seed(12345)

if (!file.exists("output")) {
  dir.create("output")
}

inputCNAqc <- function(id, dir_vcf, dir_pur, dir_cn) {
  vcf_name <- str_c(id, ".caveman_strelka2_filtered.consensus_snv.vcf.gz")
  cn_name <- str_c(id, ".battenberg.summary.csv")
  purity_name <- str_c(id, ".battenberg.rho_and_psi.txt")

  vcf <- read.table(str_c(dir_vcf, vcf_name), quote = "\"", header = F) %>% mutate(chr = V1, from = V2, to = V2, ref = V4, alt = V5, DP = 100, NV = 10, VAF = 0.1, .keep = "none")
  purity <- read.table(str_c(dir_pur, purity_name))[2, 1]
  cn <- read.csv(str_c(dir_cn, cn_name), header = F, col.names = c("X", "chromosome", "start", "end", "totalCNnorm", "minCNnorm", "totalCNtum", "minCNtum")) %>%
    filter(end - start > 1000000, chromosome != "chrX") %>%
    mutate(chr = chromosome, from = start, to = end, Major = totalCNtum - minCNtum, minor = minCNtum) %>%
    dplyr::select(chr, from, to, Major, minor)

  cnaFormat <- CNAqc::init(
    mutations = vcf,
    cna = cn,
    purity = purity,
    ref = "GRCh38"
  )
}

# CN clusters generated in Figure_7c.R
CN_clusters <- read.csv("output/CN_clusters.csv")

# Directories for VCF files, purity, and CN data
dir_vcf <- "/data/mutographs/files/Manuscripts/HNC/Somatic/SNVs/caveman_strelka2_consensus_snv_filtered/"
dir_pur <- "/data/mutographs/files/Manuscripts/HNC/Somatic/battenberg/rho_and_psi/"
dir_cn <- "/data/mutographs/files/Manuscripts/HNC/Somatic/battenberg/summary/"
list_ids <- unique(str_extract(list.files(dir_cn)[str_detect(list.files(dir_cn), "battenberg")], "PD[0-9]{5}a"))

cn_all <- NULL
for (i in 1:NROW(list_ids)) {
  cn_all[[i]] <- inputCNAqc(list_ids[[i]], dir_vcf, dir_pur, dir_cn)
}

# Genome-wide segments showing major and minor allele counts in 10 randomly picked samples per CN cluster.

for (cnv in unique(CN_clusters$HNC_clusters)) {
  cn_subset <- CN_clusters %>%
    tibble::rownames_to_column("col") %>%
    filter(HNC_clusters == cnv)
  p <- plot_multisample_CNA(cn_all[as.numeric(sample(cn_subset$col, 10))], layout = "circular") + ggplot2::labs(subtitle = cnv)
  pdf(
    file = paste0("output/Extended_Data_Figure_8a", cnv, "_", Sys.Date(), ".pdf"),
    width = 8, height = 8
  )
  print(p + theme(axis.text = element_text(size = 16)))
  dev.off()
}
